import { useEffect, useRef } from "react"
import { getIrisData, IRIS_CLASSES } from "./data";
import * as tfvis from "@tensorflow/tfjs-vis";
import * as tf from "@tensorflow/tfjs";

export default function () {
    const formDom = useRef();
    const modelRef = useRef();
    async function start() {
        const [xTrain, yTrain, xTest, yTest] = getIrisData(0.15);
        const model = tf.sequential();
        model.add(tf.layers.dense({
            units: 10,
            inputShape: [xTrain.shape[1]],
            activation: "sigmoid"
        }))
        model.add(tf.layers.dense({
            units: 3,
            activation: "softmax",
        }));
        model.compile({
            loss: "categoricalCrossentropy",
            optimizer: tf.train.adam(0.1),
            metrics: ["accuracy"]
        })
        await model.fit(xTrain, yTrain, {
            epochs: 100,
            validationData: [xTest, yTest],
            callbacks: tfvis.show.fitCallbacks(
                { name: "训练效果" },
                ["loss", "val_loss", "acc", "val_acc"],
                { callbacks: ['onEpochEnd'] }
            )
        })
        modelRef.current = model;

    }
    function predict() {
        const form = formDom.current;
        const model = modelRef.current;
        console.log(model)
        const pred = model.predict(tf.tensor([[form.a.value * 1, form.b.value * 1, form.c.value * 1, form.d.value * 1]]));
        console.log(`预测结果: ${IRIS_CLASSES[pred.argMax(1).dataSync()[0]]}`);
    }
    useEffect(() => {
        start();
    }, []);
    return <div>
        <form ref={formDom}>
            花萼长度: <input type="text" name="a" />
            花萼宽度: <input type="text" name="b" />
            花瓣长度: <input type="text" name="c" />
            花瓣宽度: <input type="text" name="d" />
        </form>
        <button onClick={predict} >预测</button>
    </div>
}